from typing import Dict, List, Set, Optional, Any
from datetime import datetime, timedelta
import asyncio
import logging
from dataclasses import dataclass

@dataclass
class WorkerStatus:
    """Status information for a worker node."""
    worker_id: str
    last_heartbeat: datetime
    current_task: Optional[str]
    memory_usage: float
    cache_usage: float
    is_active: bool

@dataclass
class TaskProgress:
    """Progress information for a task."""
    task_id: str
    layer_id: str
    status: str
    worker_id: Optional[str]
    start_time: Optional[datetime]
    end_time: Optional[datetime]
    error: Optional[str]

class RuntimeMonitor:
    """Monitors the runtime status of the distributed system.
    
    This class tracks worker health, task progress, and system resources,
    providing real-time monitoring and alerting capabilities.
    """
    
    def __init__(self, config: Dict[str, Any]):
        """Initialize the runtime monitor.
        
        Args:
            config: Configuration parameters
        """
        self.config = config
        self.logger = logging.getLogger("RuntimeMonitor")
        
        # Initialize monitoring state
        self.workers: Dict[str, WorkerStatus] = {}
        self.tasks: Dict[str, TaskProgress] = {}
        self.worker_groups: Dict[str, Set[str]] = {}
        
        # Monitoring parameters
        self.heartbeat_timeout = timedelta(seconds=config.get('heartbeat_timeout', 30))
        self.monitoring_interval = config.get('monitoring_interval', 5)
        
        # Synchronization state
        self.sync_points: Dict[int, Set[str]] = {}  # sync_point -> set of ready workers
        
        # Performance metrics
        self.system_metrics: Dict[str, List[float]] = {
            'cpu_usage': [],
            'memory_usage': [],
            'network_throughput': []
        }
    
    async def start_monitoring(self) -> None:
        """Start the monitoring loop."""
        while True:
            await self._check_worker_health()
            await self._update_system_metrics()
            await asyncio.sleep(self.monitoring_interval)
    
    def register_worker(self, worker_id: str, capabilities: Dict[str, Any]) -> None:
        """Register a new worker node.
        
        Args:
            worker_id: Worker identifier
            capabilities: Worker capabilities (GPU, memory, etc.)
        """
        self.workers[worker_id] = WorkerStatus(
            worker_id=worker_id,
            last_heartbeat=datetime.now(),
            current_task=None,
            memory_usage=0.0,
            cache_usage=0.0,
            is_active=True
        )
        self.logger.info(f"Worker {worker_id} registered with capabilities: {capabilities}")
    
    def update_worker_status(self, worker_id: str,
                            metrics: Dict[str, Any]) -> None:
        """Update status information for a worker.
        
        Args:
            worker_id: Worker identifier
            metrics: Current metrics from the worker
        """
        if worker_id not in self.workers:
            self.logger.warning(f"Received metrics from unknown worker {worker_id}")
            return
            
        worker = self.workers[worker_id]
        worker.last_heartbeat = datetime.now()
        worker.current_task = metrics.get('current_task')
        worker.memory_usage = metrics.get('memory_usage', {}).get('allocated', 0.0)
        worker.cache_usage = metrics.get('cache_stats', {}).get('usage_percent', 0.0)
    
    def register_task(self, task_id: str, layer_id: str) -> None:
        """Register a new task for monitoring.
        
        Args:
            task_id: Task identifier
            layer_id: Layer identifier
        """
        self.tasks[task_id] = TaskProgress(
            task_id=task_id,
            layer_id=layer_id,
            status='pending',
            worker_id=None,
            start_time=None,
            end_time=None,
            error=None
        )
    
    def update_task_status(self, task_id: str, status: str,
                          worker_id: Optional[str] = None,
                          error: Optional[str] = None) -> None:
        """Update the status of a task.
        
        Args:
            task_id: Task identifier
            status: New status
            worker_id: Worker processing the task
            error: Error message if failed
        """
        if task_id not in self.tasks:
            self.logger.warning(f"Received update for unknown task {task_id}")
            return
            
        task = self.tasks[task_id]
        task.status = status
        
        if status == 'running':
            task.worker_id = worker_id
            task.start_time = datetime.now()
        elif status in ['completed', 'failed']:
            task.end_time = datetime.now()
            task.error = error
    
    def register_sync_point(self, sync_point: int,
                           worker_group: List[str]) -> None:
        """Register a synchronization point for a group of workers.
        
        Args:
            sync_point: Synchronization point number
            worker_group: List of workers that need to synchronize
        """
        self.sync_points[sync_point] = set()
        self.worker_groups[str(sync_point)] = set(worker_group)
    
    def worker_ready_for_sync(self, worker_id: str,
                             sync_point: int) -> bool:
        """Mark a worker as ready for synchronization.
        
        Args:
            worker_id: Worker identifier
            sync_point: Synchronization point number
            
        Returns:
            bool: True if all workers in the group are ready
        """
        if sync_point not in self.sync_points:
            return False
            
        self.sync_points[sync_point].add(worker_id)
        worker_group = self.worker_groups.get(str(sync_point), set())
        
        return len(self.sync_points[sync_point]) == len(worker_group)
    
    async def _check_worker_health(self) -> None:
        """Check the health status of all workers."""
        current_time = datetime.now()
        
        for worker_id, worker in self.workers.items():
            if (current_time - worker.last_heartbeat > self.heartbeat_timeout and
                worker.is_active):
                worker.is_active = False
                self.logger.warning(f"Worker {worker_id} appears to be offline")
                
                # Handle any active tasks on this worker
                if worker.current_task:
                    self.update_task_status(
                        task_id=worker.current_task,
                        status='failed',
                        error='Worker disconnected'
                    )
    
    async def _update_system_metrics(self) -> None:
        """Update system-wide performance metrics."""
        # Calculate system-wide metrics
        total_memory = sum(w.memory_usage for w in self.workers.values())
        total_cache = sum(w.cache_usage for w in self.workers.values())
        active_workers = sum(1 for w in self.workers.values() if w.is_active)
        
        # Update metrics history
        self.system_metrics['memory_usage'].append(total_memory)
        if len(self.system_metrics['memory_usage']) > 100:
            self.system_metrics['memory_usage'].pop(0)
        
        # Log system status
        self.logger.info(
            f"System Status: Active Workers: {active_workers}, "
            f"Total Memory Usage: {total_memory:.2f}MB, "
            f"Total Cache Usage: {total_cache:.2f}%"
        )
    
    def get_system_status(self) -> Dict[str, Any]:
        """Get current system status.
        
        Returns:
            Dict[str, Any]: System status information
        """
        active_workers = [w for w in self.workers.values() if w.is_active]
        completed_tasks = [t for t in self.tasks.values()
                         if t.status == 'completed']
        failed_tasks = [t for t in self.tasks.values()
                       if t.status == 'failed']
        
        return {
            'workers': {
                'total': len(self.workers),
                'active': len(active_workers),
                'status': [
                    {
                        'worker_id': w.worker_id,
                        'active': w.is_active,
                        'current_task': w.current_task,
                        'memory_usage': w.memory_usage,
                        'cache_usage': w.cache_usage
                    }
                    for w in self.workers.values()
                ]
            },
            'tasks': {
                'total': len(self.tasks),
                'completed': len(completed_tasks),
                'failed': len(failed_tasks),
                'active': len([t for t in self.tasks.values()
                             if t.status == 'running'])
            },
            'system_metrics': {
                'memory_usage': self.system_metrics['memory_usage'][-1]
                if self.system_metrics['memory_usage'] else 0.0
            }
        }